"""
Joint model for relation classification.
"""
import copy
import math
import pdb

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
from torch.nn.utils import weight_norm

class LayerNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
        """Construct a layernorm module in the TF style (epsilon inside the square root).
        """
        super(LayerNorm, self).__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.bias = nn.Parameter(torch.zeros(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x):
        u = x.mean(-1, keepdim=True)
        s = (x - u).pow(2).mean(-1, keepdim=True)
        x = (x - u) / torch.sqrt(s + self.variance_epsilon)
        return self.weight * x + self.bias

class selfAttention(nn.Module):
    def __init__(self, args):
        super(selfAttention, self).__init__()
        self.num_attention_heads = args.num_attention_heads
        self.all_hidden_size = args.hidden_size
        if self.all_hidden_size % self.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (self.all_hidden_size, self.num_attention_heads))
        self.attention_head_size = int(self.all_hidden_size/self.num_attention_heads)


        self.input_size = 2*args.dim_bilstm_hidden

        self.query = nn.Linear(self.input_size, self.all_hidden_size)
        self.key = nn.Linear(self.input_size, self.all_hidden_size)
        self.value = nn.Linear(self.input_size, self.all_hidden_size)

        self.attn_dropout = nn.Dropout(args.dropout_rate)

        #all concat layer
        self.dense = nn.Linear(args.hidden_size, args.hidden_size)
        self.LayerNorm = LayerNorm(args.hidden_size, eps=1e-12)
        self.out_dropout = nn.Dropout(args.dropout_rate)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        if len(new_x_shape) == 4:
            return x.permute(0, 2, 1, 3)
        else:
            return x.permute(0, 1, 3, 2, 4)

    def forward(self, inputs):

        mixed_query_layer = self.query(inputs)
        mixed_key_layer = self.key(inputs)
        mixed_value_layer = self.value(inputs)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)

        attention_probs = nn.Softmax(dim=-1)(attention_scores)
        attention_probs = self.attn_dropout(attention_probs)
        context_layer = torch.matmul(attention_probs, value_layer)

        if len(context_layer.size()) == 4:
            context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        else:
            context_layer = context_layer.permute(0, 1, 3, 2, 4).contiguous()

        new_context_layer_shape = context_layer.size()[:-2] + (self.all_hidden_size, )
        context_layer = context_layer.view(*new_context_layer_shape)
        hidden_states = self.dense(context_layer)
        hidden_states = self.out_dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states+inputs)

        return hidden_states





class charEmbedding(nn.Module):
    '''
    Input: (max_len, max_word_len) max_len是指句子的最大长度，也就是char的batch size
    Output: (max_len, filter_num)
    '''
    def __init__(self, args):
        super(charEmbedding, self).__init__()
        self.args = args
        self.emb = nn.Embedding(self.args.size_c, self.args.dim_c)
        self.conv = weight_norm(nn.Conv1d(in_channels=self.args.dim_c, out_channels = self.args.filter_number, kernel_size=self.args.kernel_size))
        self.pool = torch.nn.MaxPool1d(self.args.max_word_len - self.args.kernel_size + 1, stride=1)
        self.drop = torch.nn.Dropout(args.dropout_rate)
        self.init()
    def init(self):
        nn.init.kaiming_uniform_(self.emb.weight.data)
        nn.init.kaiming_uniform_(self.conv.weight.data)
    def forward(self, x):
        '''
            x: one char sequence. shape: (max_len, max_word_len)
        '''
        # 如果输入是一句话

        inp = self.drop(self.emb(x))  # (max_len, max_word_len) -> (max_len, max_word_len, hidden)
        inp = inp.permute(0,2,1) # (max_len, max_word_len, hidden) -> (max_len,  hidden, max_word_len)
        out = self.conv(inp) # out: (max_len, filter_num, max_word_len - kernel_size + 1)

        return self.pool(out).squeeze() # out: (max_len, filter_num)

# 0: tokens, 1: pos, 2: mask_s, 3: labels
class Toy_model(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.embedding_matrix = args.embedding_matrix
        self.rounds = args.rounds
        self.encoder = LSTMRelationModel(args)

        # create embedding layers
        self.emb = nn.Embedding(args.vocab_size, args.dim_w, padding_idx=0)
        self.init_embeddings()
        self.emb.weight.requires_grad = True 
        self.pos_emb = nn.Embedding(len(args.pos2id), args.dim_pos, padding_idx=0)

        #char embedding
        #if self.args.use_char:
        self.char_encode = charEmbedding(args)
        self.cr_self_atten_encode = selfAttention(args)
        self.ce_self_atten_encode = selfAttention(args)

        # dropout
        self.input_dropout = nn.Dropout(args.dropout_rate)

        # GRU
        self.GRUg = torch.nn.GRUCell(5, args.dim_bilstm_hidden*2)
        self.GRUr = torch.nn.GRUCell(len(args.label2id), args.dim_bilstm_hidden*2)

        # classifer
        self.Lr = nn.Linear(4*args.dim_bilstm_hidden, 2*args.dim_bilstm_hidden)
        self.Cr = nn.Linear(2*args.dim_bilstm_hidden, len(args.label2id))
        self.Cg = nn.Linear(2*args.dim_bilstm_hidden, 5)

        # loss function 
        self.ce_loss  = nn.CrossEntropyLoss(reduction='none')
        self.bce_loss = nn.BCELoss(reduction='none')

    def init_embeddings(self):
        if self.embedding_matrix is not None:
            self.embedding_matrix = torch.from_numpy(self.embedding_matrix)
            self.emb.weight.data.copy_(self.embedding_matrix)

    # 0: tokens, 1: pos, 2: mask_s, 3: NER_labels, 4: RC_labels
    def forward(self, inputs):
        
        # Bilstm encoder
        tokens, pos, chars, mask_s, NER_labels, RC_labels = inputs

        tokens_embs = self.emb(tokens)
        rnn_inputs = [tokens_embs]

        if chars is not None:
            char_embedding = []
            for i in range(chars.shape[0]):

                one_word_char_emb = self.char_encode(chars[i])
                char_embedding.append(one_word_char_emb)

            #pdb.set_trace()
            char_embedding = torch.stack(char_embedding)
            #pdb.set_trace()
            rnn_inputs += [char_embedding]

        #rnn_inputs += [self.pos_emb(pos)]
        #pdb.set_trace()
        rnn_inputs = torch.cat(rnn_inputs, dim=2)
        lens = mask_s.sum(dim=1)
        rnn_inputs = self.input_dropout(rnn_inputs)
        H = self.encoder((rnn_inputs, lens.cpu())) 

        # generate mask 
        mask_tmp = mask_s.unsqueeze(1).repeat(1, H.size(1), 1)
        mask_RC  = torch.zeros_like(mask_s.unsqueeze(1).repeat(1, H.size(1), 1))
        len_int  = lens.int()
        for i in range(H.size(0)):
            mask_RC[i, :len_int[i], :len_int[i]] = mask_tmp[i, :len_int[i], :len_int[i]]
        
        # recurrent interaction
        Hg = H
        Hr = H
        for i in range(self.rounds):

            # Cg
            logits_Cg = self.Cg(Hg)
            prob_Cg = F.softmax(logits_Cg, dim=2)

            # Cr 
            e1  = Hr.unsqueeze(2).repeat(1, 1, H.size(1), 1)
            e2  = Hr.unsqueeze(1).repeat(1, H.size(1), 1, 1)
            e12 = torch.cat([e1, e2], dim=3)
            e12 = F.relu(self.Lr(e12), inplace=True)

            prob_Cr = torch.sigmoid(self.Cr(e12))
            prob_Cr = prob_Cr * mask_RC.unsqueeze(-1)
            prob_Cr = prob_Cr.max(dim=2)[0]
        
            # update
            Hg = self.GRUg(prob_Cg.reshape(-1, 5), H.reshape(-1, self.args.dim_bilstm_hidden*2))
            Hr = self.GRUr(prob_Cr.reshape(-1, len(self.args.label2id)), H.reshape(-1, self.args.dim_bilstm_hidden*2))
            Hg = Hg.view(H.size(0), H.size(1), self.args.dim_bilstm_hidden*2)
            Hr = Hr.view(H.size(0), H.size(1), self.args.dim_bilstm_hidden*2)
            H  = H+Hr+Hg

        # use last Hg for classification
        ## ner_attention：针对Hg设计注意力机制


        context_Hg = self.ce_self_atten_encode(Hg)

        logits_Cg = self.Cg(context_Hg)
        # pred and loss
        ner_pred = torch.argmax(logits_Cg, dim=2)
        ner_pred = ner_pred * mask_s.long()
        ner_loss = self.ce_loss(logits_Cg.reshape(-1, logits_Cg.size(-1)), NER_labels.reshape(-1))
        ner_loss = ner_loss.reshape(logits_Cg.size(0), logits_Cg.size(1))
        ner_loss = (ner_loss * mask_s).sum() / ner_loss.size(0)

        # use last Hr for classification 
        e1 = Hr.unsqueeze(2).repeat(1, 1, H.size(1), 1)
        e2 = Hr.unsqueeze(1).repeat(1, H.size(1), 1, 1)
        e12 = torch.cat([e1, e2], dim=3)
        e12 = F.relu(self.Lr(e12), inplace=True)

        # use e12
        ## cr_attention：针对Hr设计注意力机制
        #pdb.set_trace()
        context_e12 = self.cr_self_atten_encode(e12)

        logits_Cr = torch.sigmoid(self.Cr(context_e12))

        # pred and loss
        rc_pred = logits_Cr
        rc_pred = rc_pred * mask_RC.unsqueeze(-1)

        try:
            rc_loss = self.bce_loss(logits_Cr.view(-1, len(self.args.label2id)), RC_labels.view(-1, len(self.args.label2id)))
        except:
            pdb.set_trace()
        rc_loss = rc_loss.reshape(logits_Cr.size(0), logits_Cr.size(1), logits_Cr.size(2), -1)
        rc_loss = rc_loss.sum(dim=3)
        rc_loss = (rc_loss * mask_RC).sum() / rc_loss.size(0)

        return ner_loss, rc_loss, ner_pred, rc_pred

# BiLSTM model 
class LSTMRelationModel(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        #self.in_dim = args.dim_w + args.dim_pos + args.filter_number
        #pdb.set_trace()
        self.in_dim = args.dim_w   #中文按字符进行命名实体抽取，先不算pos和filter
        self.rnn = nn.LSTM(self.in_dim, args.dim_bilstm_hidden, 1, batch_first=True, dropout=0.0, bidirectional=True)

    def encode_with_rnn(self, rnn_inputs, seq_lens, batch_size):
        #pdb.set_trace()
        h0, c0 = rnn_zero_state(batch_size, self.args.dim_bilstm_hidden, 1, True)
        rnn_inputs = nn.utils.rnn.pack_padded_sequence(rnn_inputs, seq_lens, batch_first=True)
        rnn_outputs, (ht, ct) = self.rnn(rnn_inputs, (h0, c0))
        rnn_outputs, _ = nn.utils.rnn.pad_packed_sequence(rnn_outputs, batch_first=True)
        return rnn_outputs

    def forward(self, inputs):
        # unpack inputs
        inputs, lens = inputs[0], inputs[1]
        return self.encode_with_rnn(inputs, lens, inputs.size()[0])

# Initialize zero state
def rnn_zero_state(batch_size, hidden_dim, num_layers, bidirectional=True, use_cuda=True):
    total_layers = num_layers * 2 if bidirectional else num_layers
    state_shape = (total_layers, batch_size, hidden_dim)
    h0 = c0 = Variable(torch.zeros(*state_shape), requires_grad=False)
    if use_cuda:
        return h0.cuda(), c0.cuda()
    else:
        return h0, c0




